package net.cyue.ort.llm;

import java.util.List;
import java.util.function.Consumer;

// 流式输出的状态类
public class StreamingState {
    private final StringBuilder buffer;
    private final Consumer<String> textCallback;
    private final Consumer<List<Long>> tokenCallback;
    private final Runnable finishCallback;
    private final boolean shouldBuffer; // 是否需要缓冲输出
    private int lastProcessedPos; // 上次处理的位置

    public StreamingState(
        Consumer<String> textCallback,
        Consumer<List<Long>> tokenCallback,
        Runnable finishCallback,
        boolean shouldBuffer
    ) {
        this.buffer = new StringBuilder();
        this.textCallback = textCallback;
        this.tokenCallback = tokenCallback;
        this.finishCallback = finishCallback;
        this.shouldBuffer = shouldBuffer;
        this.lastProcessedPos = 0;
    }

    public void appendText(String text) {
        if (textCallback != null) {
            if (shouldBuffer) {
                buffer.append(text);
                // 尝试处理完整的字符
                processBuffer();
            } else {
                textCallback.accept(text);
            }
        }
    }
    public void appendToken(Long token, String text) {
        if (tokenCallback != null) {
            tokenCallback.accept(List.of(token));
        }

        if (textCallback != null) {
            if (shouldBuffer) {
                buffer.append(text);
                // 尝试处理完整的字符
                processBuffer();
            } else {
                textCallback.accept(text);
            }
        }
    }

    private void processBuffer() {
        String content = buffer.toString();
        int len = content.length();

        // 查找下一个完整字符的边界
        while (lastProcessedPos < len) {
            int charLen = getNextCharLength(content, lastProcessedPos);
            if (lastProcessedPos + charLen > len) {
                break;
            }

            String completeChar = content.substring(lastProcessedPos, lastProcessedPos + charLen);
            textCallback.accept(completeChar);
            lastProcessedPos += charLen;
        }
    }

    private int getNextCharLength(String content, int start) {
        // 处理UTF-8编码的字符
        char firstChar = content.charAt(start);
        if (Character.isHighSurrogate(firstChar)) {
            return 2;  // 代理对字符
        }
        return 1;  // 普通字符
    }

    public void finish() {
        if (shouldBuffer && textCallback != null) {
            // 输出剩余的缓冲内容
            String remaining = buffer.substring(lastProcessedPos);
            if (!remaining.isEmpty()) {
                textCallback.accept(remaining);
            }
        }

        if (finishCallback != null) {
            finishCallback.run();
        }
    }
}
